import lithosim_cuda as ilt
import torch
import torchvision

if __name__ == '__main__':
    mask = ilt.load_image('../ICCAD2013/png/target1.png')
    threshold = 0.225
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    kernel_focus = torch.load('../../kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)  # 24, 35, 35
    kernel_focus_ct = torch.load('../../kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)  # 24, 35, 35
    weight_focus = torch.load('../../kernel_neuralilt/weight_focus_tensor.pt', map_location=device)  # 24
    save_bin_wafer_image = True
    wafer_output_path = '../neur_wafer1.png'
    intensity_output_path = '../neur_intensity1.png'

    intensity_map, binary_wafer = ilt.lithosim(mask, threshold, kernel_focus, weight_focus,
                                                       wafer_output_path, save_bin_wafer_image)

    intensity_2d = intensity_map.squeeze(0)
    torchvision.utils.save_image(intensity_2d, intensity_output_path)
    print("Save binary wafer image in %s" % intensity_output_path)
